/* Copyright (C) 2021 Rain */

/* This file is part of XNIX. */

/* 
  XNIX is free software: you can redistribute it and/or modify 
  it under the terms of the GNU General Public License as published by 
  the Free Software Foundation, either version 3 of the License, or 
  (at your option) and later version. 
*/

/*
  XNIX is distributed in the hope that it will be useful, 
  but WITHOUT ANY WARRANTY; without even the implied warranty of 
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 
  GNU General Public License for more details. 
*/

/*
  You should have received a copy of the GNU General Public License 
   along with XNIX. If not, see <https://www.gnu.org/licenses/>.  
*/




#include <kernel/init.h>
#include <kernel/errno.h>
#include <kernel/types.h>

#include <kernel/panic.h>
#include <kernel/int.h>
#include <kernel/sched.h>

#include <kernel/console.h>
#include <kernel/mm.h>

#include <sys/syscall.h>

#include <arch/x86/int.h>
#include <arch/io.h>
#include <arch/sys.h>

#include <arch/x86/tss.h>
#include <arch/x86/paging.h>

#include <lib/string.h>



/* kernel/trap.S */

extern void ret_intr(void);

struct task init_task;
struct task *current;

struct task *tasks[MAX_TASKS] = {0};


int do_exit(__uint64_t exitcode) {
	printk(white, "do_exit: %q\n", exitcode);

	for (;;);
}


int task_init(void) {
	/* initial the init-tss */

	init_tss.rsp0 = init_tss.rsp1 = init_tss.rsp2 = get_rsp();
	init_tss.ist1 = init_tss.ist2 = init_tss.ist3 = get_rsp();
	init_tss.ist4 = init_tss.ist5 = init_tss.ist6 = get_rsp();
	init_tss.ist7 = get_rsp();

	set_tss(&init_tss, 0);

	current = &init_task;

	return 0;
}


/* 
 * entry for all kernel tasks. (execute only one time)
 * get registers from stack (store in do_fork, memcpy). 
 * and call function %rcx with argument %rdi, then call 
 * do_exit. 
 * */

__asm__(
	"ktask_entry:\n\t"
		"popq %rax\n\t"
		"movq %rax, %gs\n\t"
		"popq %rax\n\t"
		"movq %rax, %fs\n\t"
		"popq %rax\n\t"
		"movq %rax, %es\n\t"
		"popq %rax\n\t"
		"movq %rax, %ds\n\t"

		"popq %r15\n\t"
		"popq %r14\n\t"
		"popq %r13\n\t"
		"popq %r12\n\t"

		"popq %r11\n\t"
		"popq %r10\n\t"
		"popq %r9\n\t"
		"popq %r8\n\t"

		"popq %rdi\n\t"
		"popq %rsi\n\t"
		"popq %rbp\n\t"

		"popq %rbx\n\t"
		"popq %rdx\n\t"
		"popq %rcx\n\t"
		"popq %rax\n\t"

		"addq $16, %rsp\n\t"
		"popfq\n\t"
		"addq $16, %rsp\n\t"
		"sti\n\t"

		"callq *%rcx\n\t"

		/* syscall 0: do_exit */
		"movq %rax, %rbx\n\t"
		"movq $0, %rax\n\t"

		"int $0x80\n\t"
		/* it never here */
);


/* 
 * create a kernel task. 
 * (just set some registers)
 * */

int ktask_create(__uint64_t (* entry) (__uint64_t), __uint64_t arg) {
	struct registers regs;

	memset(&regs, 0x00, sizeof(regs));

	/* entry function */
	regs.rcx = (__uint64_t) entry;
	
	/* argument */
	regs.rdi = arg;

	return do_fork(&regs, 0);
}

int utask_create(__uint64_t (* entry) (__uint64_t), __uint64_t arg) {
	struct registers regs;

	memset(&regs, 0x00, sizeof(regs));

	regs.rip = (__uint64_t) entry;
	regs.rdi = arg;

	return do_fork(&regs, 3);
}



static __uint32_t find_free_pid(void) {
	__uint32_t freepid;

	for (freepid = 0; freepid < MAX_TASKS; freepid++) 
		if (!tasks[freepid]) 
			return freepid;

	return -1;
}



/* 
 * we use this function to create a new task. 
 *
 * it doesn't like that function that we often use in linux (no argument), 
 * it needs some arguments. 
 *
 * if now is a user task, its child should be a user task too. 
 *
 * we store reg on the top of stack, because ret_intr (or ktask_entry) will 
 * popq all these registers when the task runs. 
 * */

int do_fork(struct registers *regs, __uint64_t dpl) {
	__uint64_t start_stack;

	if (!dpl) 
		start_stack = (__uint64_t) kmalloc(&pool, 65536);

	else 
		start_stack = (__uint64_t) kmalloc(&upool, 65536);

	if (!start_stack) 
		return -ENOMEM;

	/* += stack_size: point to the top of stack. */

	start_stack += 65536;

	struct task *child = kmalloc(&pool, sizeof(struct task));

	if (!child) 
		return -ENOMEM;

	/* current is this task's father */
	memcpy(child, current, sizeof(struct task));

	child->state = STATE_IDE;

	child->utime = child->ktime = 0;

	child->start_time = jiffies;

	child->signal = 0;

	/* allow all signals */
	child->signal_mask = 0xffffffff;

	child->exit_code = 0;

	child->pid = find_free_pid();

	if (child->pid == -1) 
		return -EAGAIN;

	printk(white, "new task pid: %d\n", child->pid);

	child->father = current->pid;

	child->cr3 = 0x0000000000000000;

	/* kernel task */
	if (!dpl) {
		__asm__ __volatile__ (
			"leaq ktask_entry(%%rip), %%rax\n\t"
			"movq %%rax, %0\n\t"

			: "=m" (child->rip)
			:: "ax"
		);
		
		/* I don't know why this doesn't work. . . */

//		child->rip = (__uint64_t) &ktask_entry;

		child->dataseg = KERNEL_DS;
		child->codeseg = KERNEL_CS;

		/* bullshit! -> ktask_entry will popq it in the stack. */
		/* we don't need to set ss in regs, because
		 * ktask_entry will ignore them
		 * */
		regs->gs = regs->fs = regs->es = regs->ds = KERNEL_DS;
	}

	/* dpl = 3 */
	else {
		__asm__ __volatile__ (
			"leaq ret_intr(%%rip), %%rax\n\t"
			"movq %%rax, %0\n\t"

			: "=m" (child->rip)
			:: "ax"
		);

//		child->rip = (__uint64_t) &ret_intr;
		child->dataseg = USER_DS;
		child->codeseg = USER_CS;

		regs->ss = regs->fs = regs->es = regs->ds = regs->gs = USER_DS;
		regs->cs = USER_CS;
		regs->rflags = 1 << 9;
		/* set from the utask_create */
//		regs->rip = ;
		regs->rsp = start_stack - sizeof(struct registers);
	}

	/* 20 ms, default */
	child->counter = 2;

	child->nice = 2;
	child->rsp = start_stack - sizeof(struct registers);

	__asm__ __volatile__ (
		"cld;rep;movsb;\n\t"
		:: "S" (regs), "D" (child->rsp), "c" (sizeof(struct registers))
	);

	child->state = STATE_READY;

	tasks[child->pid] = child;

	return 0;
}


void set_tss(struct tss *table, __uint64_t dpl) {
	__uint64_t tss = (__uint64_t) table;

	/* create a descriptor (in gdt) and register it */

	struct tss_entry *entry = (struct tss_entry *) (gdt_table + 5);

	entry->low = entry->high = 0x0000000000000000;

	/* `103` is the size of a tss table */

	entry->low = 103ul | ((tss & 0xffff) << 16) | (((tss >> 16) & 0xff) << 32) 
		| ((0x80 + dpl + 0x09) << 40) | (tss << 56 & 0xff);

	entry->high = (tss >> 32) & 0xffffffff;

	/* load it to tr */
	ltr(40);

	return;
}

/* well... but this is ok */
static char need_open_clock = 0;

void schedule(void) {
	/* find and switch to a new task */

	__uint32_t pid;
	__uint32_t best = 17;
	__uint32_t bestpid = 0;

	for (pid = 0; pid < MAX_TASKS; pid++) {
		if (tasks[pid] && tasks[pid]->state == STATE_READY 
				&& tasks[pid]->nice < best) {
			bestpid = pid;
			best = tasks[pid]->nice;
		}
	}

	if (best == 17) {
		/* there must are some STATE_FINISH, change them to STATE_READY 
		 * and get bestpid again. 
		 * */
		
		int i;
		best = 17;

		for (i = 0; i < MAX_TASKS; i++) {
			if (tasks[i] && tasks[i]->state == STATE_FINISH) {
				tasks[i]->state = STATE_READY;
				tasks[i]->counter = tasks[i]->nice;

				if (tasks[i]->nice < best) {
					bestpid = i;
					best = tasks[pid]->nice;
				}
			}
		}
	}
	
	if (current->pid == bestpid) 
		return;

	/* will sti() in switch_to() */

	switch_to(tasks[bestpid]);

	return;
}

void do_clock(void) {
	/* still in the init mode */

	if (!current) 
		return;

	current->counter--;

	if (!current->counter) {
		current->state = STATE_FINISH;
		need_open_clock = 1;

		schedule();
	}

	return;
}
